# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import datetime


def expiry_accurate_datetime(expiry) -> datetime.datetime:
  assert isinstance(expiry, datetime.datetime)
  assert expiry.tzinfo is None
  # Settle time is utc 8am on Friday.
  assert expiry.weekday() == 4
  expiry = expiry.replace(hour=8)
  return expiry


def find_this_week_contract_expiring_date(current_datetime):
  expiry = current_datetime
  # Settle time is utc 8am on Friday.
  if expiry.weekday() == 4 and expiry.hour >= 8:
    # Skip this week Friday because it just expired.
    expiry = expiry + datetime.timedelta(days=1)
  while expiry.weekday() != 4:
    expiry = expiry + datetime.timedelta(days=1)
  return expiry_accurate_datetime(expiry)


def find_next_week_contract_expiring_date(current_datetime):
  expiry = current_datetime
  if expiry.weekday() == 4 and expiry.hour >= 8:
    # Skip this week Friday because it just expired.
    expiry = expiry + datetime.timedelta(days=1)
  while expiry.weekday() != 4:
    expiry = expiry + datetime.timedelta(days=1)
  expiry = expiry + datetime.timedelta(days=1)
  while expiry.weekday() != 4:
    expiry = expiry + datetime.timedelta(days=1)
  return expiry_accurate_datetime(expiry)


def find_quarter_contract_expiring_date(quarter_contract_date, current_datetime):
  for date in quarter_contract_date:
    date = datetime.datetime.strptime(date, '%Y%m%d')
    date = expiry_accurate_datetime(date)
    if date - current_datetime > datetime.timedelta(days=14):
      return date


def find_next_quarter_contract_expiring_date(quarter_contract_date, current_datetime):
  quarter_date = find_quarter_contract_expiring_date(quarter_contract_date, current_datetime)
  quarter_date_str = quarter_date.strftime('%Y%m%d')
  idx = quarter_contract_date.index(quarter_date_str)
  assert idx < len(quarter_contract_date) - 1, 'Need to add new quarter contract dates!'
  next_quarter_date_str = quarter_contract_date[idx + 1]
  return expiry_accurate_datetime(datetime.datetime.strptime(next_quarter_date_str, '%Y%m%d'))
